# coding=utf-8

import json
import os
import math
import uuid
import random

# 文件结构：
# [定日镜面积]/[定日镜长宽比]/[定日镜玻璃厚度]/{unique_id}.bin
# 定日镜长宽比：长_宽，如8_5，4_3，1_1
# 玻璃厚度：单位 mm
# 叶子文件夹中存放input_args.json，保存{unique_id}.bin中的参数
# flux_map_shape: [长度，宽度] （in pixel）
# flux_map_shape：定日镜与接收器距离最近处长宽取定日镜长度的1.5倍，最远处取4倍
# peak_flux_param：定日镜单位面积上的总能量（卷积模型）
# fitted: 是否已完成参数拟合
# example:
# [
#     {
#         "id": 0,
#         "fitted": false,
#         "flux_map_shape": [160, 160]
#         "distance_heliostat_receiver": 25,
#         "sun_elevation": 90,
#         "sun_azimuth": 0,
#         "heliostat_glass_sigma": 0.002,
#         "sun_csr": 0.1,
#         "heliostat_glass_refractive_index": 1.5
#         "peak_flux_param": 6.901197509765625E2,
#         "sigmas": [
#             4.987233877182007E-2,
#             4.987233877182007E-2
#         ],
#         "UNIZAR": {
#             "x13": 2.1725692749023438E0,
#             "x10": -1.6070754528045654E0,
#             "y11": -1.3468780517578125E0,
#             "y10": 1.3304443359375E0,
#             "gradient_inv": 2.1428488194942474E-1
#         }
#     }
# ]

SOLAR_3_PATH = "D:/Learn/solar_software/solar_3/build/appsolar_2.exe"
HELIOSTAT_SIZE = 151
HELIOSTAT_L_W_RATIO = 8.0 / 5.0
GLASS_THICKNESS = 0.001
PATH_PREFIX = "./151/8_5/1/"


def lerp(__a: float, __b: float, __t: float) -> float:
    """return __a * (1 - __t) + b * __t"""
    """A copy of C++'s std::lerp from libc++"""
    if (__a <= 0 and __b >= 0) or (__a >= 0 and __b <= 0):
        return __t * __b + (1 - __t) * __a

    if __t == 1:
        return __b
    __x = __a + __t * (__b - __a)
    if (__t > 1) == (__b > __a):
        return __x if __b < __x else __b
    else:
        return __x if __x < __b else __b


def gen_unizar_init_model_params(software_path: str, input_json_path: str):
    with open(input_json_path, "r") as input_json_file:
        input_json_object = json.load(input_json_file)
        input_json_object["algorithm"] = "analytical"
        input_json_object["model_type"] = "convolution"
        input_json_object["model_name"] = "UNIZAR"

    with open(input_json_path, "w") as input_json_file:
        json.dump(input_json_object, input_json_file)

    tmp_output_json_path = str(uuid.uuid1()) + ".json"
    try:
        os.system(
            software_path
            + " --cli "
            + input_json_path
            + " -o "
            + tmp_output_json_path
            + " --dump-args"
        )
    except:
        print("gen_unizar_init_model_params failed")

    json_file = open(tmp_output_json_path, "r")
    json_object = json.load(json_file)
    json_file.close()
    os.remove(tmp_output_json_path)
    return json_object["heliostat_args"][0]


def gen_raytracing_binary_one_heliostat(
    software_path: str, input_json_path: str, output_path: str
):
    with open(input_json_path, "r") as input_json_file:
        input_json_object = json.load(input_json_file)
        input_json_object["algorithm"] = "raytracing"

    with open(input_json_path, "w") as input_json_file:
        json.dump(input_json_object, input_json_file)

    try:
        os.system(software_path + " --cli " + input_json_path + " -o " + output_path)
    except:
        print("gen_raytracing_binary_one_heliostat failed")


def gen_image_plane_size(
    heliostat_length: float, current_distance: float, max_distance: float
):
    return math.ceil(heliostat_length * lerp(1.5, 4, current_distance / max_distance))


def gen_image_plane_shape(image_plane_length: int):
    # image_plane_length / 0.05
    return (image_plane_length * 20, image_plane_length * 20)


def gen_sun_csr():
    return random.uniform(0, 0.4)


def gen_heliostat_glass_sigma():
    random_sigma = random.uniform(0, 3)
    if random.uniform(0, 3) < 2:
        random_sigma = random_sigma / 3.0
    return random_sigma / 1000.0


def gen_heliostat_glass_refractive_index():
    return random.uniform(1.4, 1.6)


def gen_result_one_heliostat(
    current_id,
    flux_map_shape,
    distance_heliostat_receiver,
    sun_elevation,
    sun_azimuth,
    heliostat_glass_sigma,
    sun_csr,
    heliostat_glass_refractive_index,
    unizar_init_model_params,
):
    result_dict = {
        "id": current_id,
        "flux_map_shape": flux_map_shape,
        "distance_heliostat_receiver": distance_heliostat_receiver,
        "sun_elevation": sun_elevation,
        "sun_azimuth": sun_azimuth,
        "heliostat_glass_sigma": heliostat_glass_sigma,
        "sun_csr": sun_csr,
        "heliostat_glass_refractive_index": heliostat_glass_refractive_index,
    }
    result_dict.update(unizar_init_model_params)
    return result_dict


def gen_temp_input_json(
    template_json_object,
    heliostat_length,
    heliostat_width,
    distance_heliostat_receiver,
    heliostat_glass_thickness,
    sun_azimuth,
    sun_elevation,
    heliostat_glass_sigma,
    sun_csr,
    heliostat_glass_refractive_index,
    rectangle_receiver_length,
    rectangle_receiver_width,
):
    template_json_object["field"]["heliostats"]["group"]["geometry"]["size"][
        "x"
    ] = heliostat_length
    template_json_object["field"]["heliostats"]["group"]["geometry"]["size"][
        "y"
    ] = heliostat_width
    template_json_object["field"]["heliostats"]["group"]["differ"][0]["position"][
        "z"
    ] = distance_heliostat_receiver
    template_json_object["field"]["heliostats"]["group"]["material"][
        "glass_thickness"
    ] = heliostat_glass_thickness
    template_json_object["config"]["sun"]["sun_position"]["azimuth"] = sun_azimuth
    template_json_object["config"]["sun"]["sun_position"]["elevation"] = sun_elevation
    template_json_object["field"]["heliostats"]["group"]["material"][
        "slope"
    ] = heliostat_glass_sigma
    template_json_object["config"]["sun"]["sun_shape"]["csr"] = sun_csr
    template_json_object["field"]["heliostats"]["group"]["material"][
        "glass_refractive_index"
    ] = heliostat_glass_refractive_index
    template_json_object["field"]["receivers"]["group"]["geometry"]["size"][
        "x"
    ] = rectangle_receiver_length
    template_json_object["field"]["receivers"]["group"]["geometry"]["size"][
        "y"
    ] = rectangle_receiver_width


def main():
    if not os.path.exists(PATH_PREFIX):
        os.makedirs(PATH_PREFIX)

    heliostat_width = math.sqrt(HELIOSTAT_SIZE / HELIOSTAT_L_W_RATIO)
    heliostat_length = heliostat_width * HELIOSTAT_L_W_RATIO
    sigma_total = 5.111 * 0.001
    max_distance = math.ceil(
        math.hypot(heliostat_length, heliostat_width) / (4 * sigma_total)
    )
    distance_step = 25
    sun_elevation_step = 15
    sun_azimuth_step = 20
    current_id = 0
    # results = []

    with open(
        "D:/Learn/solar_software/pytorch_script/train_data_template.json", "r+"
    ) as json_file:
        json_object = json.load(json_file)
        for current_distance in range(distance_step, max_distance, distance_step):
            for current_sun_elevation in range(15, 90, sun_elevation_step):
                for current_sun_azimuth in range(0, 181, sun_azimuth_step):
                    current_output_json_path = str(current_id) + ".json"
                    current_output_bin_path = str(current_id) + ".bin"
                    if (os.path.exists(PATH_PREFIX + current_output_json_path) and os.path.exists(PATH_PREFIX + current_output_bin_path)):
                        current_id = current_id + 1
                        continue

                    image_plane_length = gen_image_plane_size(
                        heliostat_length, current_distance, max_distance
                    )
                    image_plane_shape = gen_image_plane_shape(image_plane_length)
                    sun_csr = gen_sun_csr()
                    heliostat_glass_sigma = gen_heliostat_glass_sigma()
                    heliostat_glass_refractive_index = (
                        gen_heliostat_glass_refractive_index()
                    )
                    gen_temp_input_json(
                        json_object,
                        heliostat_length,
                        heliostat_width,
                        current_distance,
                        GLASS_THICKNESS,
                        current_sun_azimuth,
                        current_sun_elevation,
                        heliostat_glass_sigma,
                        sun_csr,
                        heliostat_glass_refractive_index,
                        image_plane_length,
                        image_plane_length,
                    )
                    temp_input_json_path = str(uuid.uuid1()) + ".json"
                    with open(temp_input_json_path, "w") as temp_input_json_file:
                        json.dump(json_object, temp_input_json_file)
                    temp_input_json_file.close()
                    # print(temp_input_json_path)
                    unizar_init_model_params = gen_unizar_init_model_params(
                        SOLAR_3_PATH, temp_input_json_path
                    )
                    gen_raytracing_binary_one_heliostat(
                        SOLAR_3_PATH,
                        temp_input_json_path,
                        PATH_PREFIX + str(current_id) + ".bin",
                    )
                    os.remove(temp_input_json_path)
                    result_one_heliostat = gen_result_one_heliostat(
                        current_id,
                        image_plane_shape,
                        current_distance,
                        current_sun_elevation,
                        current_sun_azimuth,
                        heliostat_glass_sigma,
                        sun_csr,
                        heliostat_glass_refractive_index,
                        unizar_init_model_params,
                    )
                    # results.append(result_one_heliostat)
                    
                    with open(PATH_PREFIX + current_output_json_path, "w") as current_output_json_file:
                        json.dump(result_one_heliostat, current_output_json_file)

                    current_id = current_id + 1

        # sun_elevation = 90
        for current_distance in range(distance_step, max_distance, distance_step):
            # for current_sun_azimuth in range(0, 181, sun_azimuth_step):
            current_output_json_path = str(current_id) + ".json"
            current_output_bin_path = str(current_id) + ".bin"
            if (os.path.exists(PATH_PREFIX + current_output_json_path) and os.path.exists(PATH_PREFIX + current_output_bin_path)):
                current_id = current_id + 1
                continue
            
            current_sun_elevation = 90
            current_sun_azimuth = 0
            image_plane_length = gen_image_plane_size(
                heliostat_length, current_distance, max_distance
            )
            image_plane_shape = gen_image_plane_shape(image_plane_length)
            sun_csr = gen_sun_csr()
            heliostat_glass_sigma = gen_heliostat_glass_sigma()
            heliostat_glass_refractive_index = gen_heliostat_glass_refractive_index()
            gen_temp_input_json(
                json_object,
                heliostat_length,
                heliostat_width,
                current_distance,
                GLASS_THICKNESS,
                current_sun_azimuth,
                current_sun_elevation,
                heliostat_glass_sigma,
                sun_csr,
                heliostat_glass_refractive_index,
                image_plane_length,
                image_plane_length,
            )
            temp_input_json_path = str(uuid.uuid1()) + ".json"
            with open(temp_input_json_path, "w") as temp_input_json_file:
                json.dump(json_object, temp_input_json_file)
            unizar_init_model_params = gen_unizar_init_model_params(
                SOLAR_3_PATH,
                temp_input_json_path,
            )
            gen_raytracing_binary_one_heliostat(
                SOLAR_3_PATH,
                temp_input_json_path,
                PATH_PREFIX + str(current_id) + ".bin",
            )
            os.remove(temp_input_json_path)
            result_one_heliostat = gen_result_one_heliostat(
                current_id,
                image_plane_shape,
                current_distance,
                current_sun_elevation,
                current_sun_azimuth,
                heliostat_glass_sigma,
                sun_csr,
                heliostat_glass_refractive_index,
                unizar_init_model_params,
            )
            # results.append(result_one_heliostat)
            current_output_json_path = str(current_id) + ".json"
            with open(PATH_PREFIX + current_output_json_path, "w") as current_output_json_file:
                json.dump(result_one_heliostat, current_output_json_file)

            current_id = current_id + 1

    # output_json_path = "params.json"
    # with open(PATH_PREFIX + output_json_path, "w") as output_json_file:
    #     json.dump(results, output_json_file)

        # algorithm_type = json_object["algorithm"]
        # heliostat_length = json_object["field"]["heliostats"]["group"]["geometry"][
        #     "size"
        # ]["x"]
        # heliostat_width = json_object["field"]["heliostats"]["group"]["geometry"][
        #     "size"
        # ]["y"]
        # distance_heliostat_receiver = json_object["field"]["heliostats"]["group"][
        #     "differ"
        # ][0]["position"]["y"]
        # heliostat_glass_thickness = json_object["field"]["heliostats"]["group"][
        #     "material"
        # ]["glass_thickness"]
        # sun_azimuth = json_object["config"]["sun"]["sun_position"]["azimuth"]
        # sun_elevation = json_object["config"]["sun"]["sun_position"]["elevation"]
        # heliostat_glass_sigma = json_object["field"]["heliostats"]["group"]["material"][
        #     "slope"
        # ]  # 由于RMCRT算法限制，玻璃表面和涂层面的sigma相同
        # sun_csr = json_object["config"]["sun"]["sun_shape"]["csr"]
        # heliostat_glass_refractive_index = json_object["field"]["heliostats"]["group"][
        #     "material"
        # ]["glass_refractive_index"]
        # rectangle_receiver_length = json_object["field"]["receivers"]["group"][
        #     "geometry"
        # ]["size"]["x"]
        # rectangle_receiver_width = json_object["field"]["receivers"]["group"][
        #     "geometry"
        # ]["size"]["y"]

        # print(
        #     algorithm_type,
        #     heliostat_length,
        #     heliostat_width,
        #     distance_heliostat_receiver,
        #     heliostat_glass_thickness,
        #     sun_azimuth,
        #     sun_elevation,
        #     heliostat_glass_sigma,
        #     sun_csr,
        #     heliostat_glass_refractive_index,
        # )


if __name__ == "__main__":
    main()
